#ifndef FSM_H
#define FSM_H

/**
 * @file fsm.h
 * @brief 在SheriffOS/FSM的基础上进行了优化：
 *
 * 单例工厂避免手动定义实例，状态绑定状态机
 * Case由模板收集，所有初始化在全局构造时完成；使用模板lambda替代std::function减小性能开销；使用静态id区别实现，避免string标签
 * 全新语法糖 & 编译时检查充分防呆
 */

#include <rtos/critical_section.h>
#include <algorithm>
#include <cstdint>
#include <memory>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
#include <core.h>

namespace os {
class State;
class StateMachine;
}  // namespace os

namespace detail {
class Branch;
class StateCollector;
class TransitionBuilder;
}  // namespace detail

namespace detail {
/**
 * @class Case
 * @brief 代表一个分支（Branch）中的一个条件转移
 * @tparam F 条件函数的类型，通常是一个lambda，模板方法避免function开销
 */
template <typename F>
class Case {
public:
    F condition_;
    os::State& toState_;

    explicit Case(F&& c, os::State& s) : condition_(std::forward<F>(c)), toState_(s) {}
};

template <typename F>
Case(F&&, os::State&) -> Case<F>;

/**
 * @class Branch
 * @brief 转移逻辑中分支的抽象基类，目的是擦除实现的模板
 */
class Branch {
public:
    virtual ~Branch() = default;

    Branch() = default;

    /**
     * @brief 按顺序判断所有Case，返回第一个满足条件的目标状态
     * @return 如果有条件满足，则为目标状态指针；否则为nullptr
     */
    [[nodiscard]] virtual os::State* get_toState() const = 0;

    /**
     * @brief 获取该分支中所有可能的目标状态
     * @return 包含所有目标状态指针的vector
     */
    [[nodiscard]] virtual std::vector<os::State*> get_branch_state() const = 0;
};

/**
 * @class BranchImpl
 * @brief Branch的实现，使用std::tuple存储异构Case对象
 * @tparam Cases Case对象的类型列表
 */
template <typename... Cases>
class BranchImpl final : public Branch {
    std::tuple<Cases...> cases_;

public:
    explicit BranchImpl(Cases&&... cases) : cases_(std::forward<Cases>(cases)...) {}

    [[nodiscard]] os::State* get_toState() const override {
        os::State* targetState = nullptr;
        std::apply(
            [&](auto const&... c) {
                auto check = [&](auto const& cs) {
                    if (cs.condition_()) {
                        targetState = &cs.toState_;
                        return true;
                    }
                    return false;
                };
                (check(c) || ...);
            },
            cases_);
        return targetState;
    }

    [[nodiscard]] std::vector<os::State*> get_branch_state() const override {
        std::vector<os::State*> vec;
        vec.reserve(sizeof...(Cases));
        std::apply([&](auto const&... c) { (vec.push_back(&c.toState_), ...); }, cases_);
        return vec;
    }
};

/**
 * @brief 为二分查找提供
 */
constexpr uint8_t linear_threshold = 16;
}  // namespace detail

namespace os {
/**
 * @brief Branch工厂
 * @param cases 构成该分支的所有Case对象
 * @return 指向新创建的Branch对象的unique_ptr
 */
template <typename... Cases>
[[nodiscard]] auto branch(Cases&&... cases) {
    return std::make_unique<detail::BranchImpl<Cases...>>(std::forward<Cases>(cases)...);
}

/**
 * @brief 事件基类
 */
class Event {};

/**
 * @brief 状态基类
 */
class State {
    friend detail::TransitionBuilder;
    friend StateMachine;
    friend detail::StateCollector;

    /**
     * @brief 存储事件发送后对应Branch信息
     */
    struct Trans {
        /**
         * @brief 使用事件对象地址作为键值
         */
        std::uintptr_t event_key_{};
        std::unique_ptr<detail::Branch> branch_;

        Trans(std::uintptr_t const key, std::unique_ptr<detail::Branch> branch)
            : event_key_(key), branch_(std::move(branch)) {}

        /**
         * @brief 为二分查找提供
         */
        bool operator<(Trans const& other) const { return event_key_ < other.event_key_; }
    };

    std::vector<Trans> transLogic_{};
    StateMachine* fsm_ = nullptr;
    bool isSorted_ = false;

protected:
    /**
     * @brief 获取所属的状态机
     * @note 允许通过fsm()在onEnter, onExecute, onExit中访问当前状态机
     */
    [[nodiscard]] StateMachine& fsm() const { return *fsm_; }

public:
    State() = default;

    virtual ~State() = default;

    // 禁止拷贝和移动，因为状态应该是唯一的单例
    State(State const&) = delete;

    State& operator=(State const&) = delete;

    State(State&&) = delete;

    State& operator=(State&&) = delete;

private:
    virtual void onEnter() {}

    virtual void onExecute() {}

    virtual void onExit() {}

    void addLogic(Event& event, std::unique_ptr<detail::Branch> branch) {
        transLogic_.emplace_back(reinterpret_cast<std::uintptr_t>(&event), std::move(branch));
        isSorted_ = false;  // 添加新逻辑后需要重新排序
    }

    /**
     * @brief 检查事件发生是是否存在转移逻辑
     * @param event 当前发生的事件
     * @return 转移的目标状态，当没有转移逻辑时返回nullptr
     */
    [[nodiscard]] State* checkLogic(Event const& event) const {
        auto const key = reinterpret_cast<std::uintptr_t>(&event);

        // 数据量小时采用线性查找
        if (transLogic_.size() < detail::linear_threshold) {
            for (auto const& t : transLogic_) {
                if (t.event_key_ == key) return t.branch_->get_toState();
            }
            return nullptr;
        }

        // 数据量大时采用二分查找优化性能
        auto const it = std::lower_bound(transLogic_.begin(), transLogic_.end(), key,
                                         [](Trans const& t, std::uintptr_t const k) { return t.event_key_ < k; });

        if (it != transLogic_.end() && it->event_key_ == key) {
            return it->branch_->get_toState();
        }

        return nullptr;
    }

    /**
     * @brief 将状态绑定到状态机
     * @param fsm 状态机
     */
    void bind(StateMachine& fsm) {
        if (fsm_ == nullptr) {
            fsm_ = &fsm;
        } else if (fsm_ != &fsm) {
            osError("State cannot be bound to multiple StateMachines.");
        }
    }

    /**
     * @brief 为二分查找排序
     */
    void sortTransitions() {
        if (!isSorted_ && transLogic_.size() >= detail::linear_threshold) {
            std::sort(transLogic_.begin(), transLogic_.end());
            isSorted_ = true;
        }
    }
};

/**
 * @brief 类型单例工厂
 */
template <typename S>
[[nodiscard]] S& state() {
    static_assert(std::is_base_of_v<State, S>, "S type must inherit from State");
    static S instance;
    return instance;
}

/**
 * @brief 状态单例工厂
 */
template <typename E>
[[nodiscard]] E& event() {
    static_assert(std::is_base_of_v<Event, E>, "E type must inherit from Event");
    static E instance;
    return instance;
}
}  // namespace os

namespace detail {
/**
 * @class StateCollector
 * @brief 一个临时的辅助类，用于在状态机初始化时收集所有参与定义的状态
 */
class StateCollector {
public:
    std::vector<os::State*> states;

    explicit StateCollector(std::vector<os::State*>&& v) : states(std::move(v)) {}

    /**
     * @brief 遍历收集到的状态并进行绑定和排序操作
     * @param f 当前状态机
     */
    void bindState(os::StateMachine& f) const {
        for (auto& s : states) {
            s->bind(f);
            s->sortTransitions();
        }
    }
};
}  // namespace detail

namespace os {
/**
 * @class StateMachine
 * @brief 状态机本体
 */
class StateMachine {
public:
    using StatePtr = State*;

    /**
     * @brief 构造函数，接收一个初始化列表来定义所有状态和转移
     * @param list 包含所有转移定义和状态声明的初始化列表
     */
    StateMachine(std::initializer_list<detail::StateCollector> const list) {
        for (auto const& i : list) {
            i.bindState(*this);
        }
    }

    explicit StateMachine() = default;

    ~StateMachine() {
        if (currentState) {
            currentState->onExit();
        }
    }

    // 禁止拷贝和移动
    StateMachine(StateMachine const&) = delete;

    StateMachine& operator=(StateMachine const&) = delete;

    StateMachine(StateMachine&&) = delete;

    StateMachine& operator=(StateMachine&&) = delete;

    /**
     * @brief 强制切换状态，也用作设置初始状态
     * @tparam S 目标状态类型
     */
    template <typename S>
    void checkoutTo() {
        static_assert(std::is_base_of_v<State, S>, "S must be a State type");
        SECTION_SAFE_FROM_TASKS {
            State* targetState = &state<S>();
            if (targetState->fsm_ != this) {
                osError("Target state is not bound to this FSM.");
            }
            if (currentState != targetState) {
                transitionTo(targetState);
            }
        }
    }

    /**
     * @brief 向状态机提交一个事件，触发可能的状态转移
     * @tparam E 事件
     */
    template <typename E>
    void react() {
        static_assert(std::is_base_of_v<Event, E>, "E must be an Event type");
        if (!currentState) {
            osError("FSM has no initial state. Call checkoutTo<S>() first.");
        }
        SECTION_SAFE_FROM_TASKS {
            if (State* to_state = currentState->checkLogic(event<E>())) {
                if (to_state->fsm_ != this) {
                    osError("Target state is not bound to this FSM.");
                }
                transitionTo(to_state);
            }
        }
    }

    /**
     * @brief 执行当前状态的onExecute
     */
    void execute() const {
        if (currentState) {
            currentState->onExecute();
        }
    }

    /**
     * @brief 检查状态机当前是否处于指定状态
     * @tparam S 状态类型
     * @return 如果当前是该状态，则为true，否则为false
     */
    template <typename S>
    [[nodiscard]] bool is_in() const {
        static_assert(std::is_base_of_v<State, S>, "S must be a State type");
        return currentState == &state<S>();
    }

private:
    StatePtr currentState = nullptr;

    void transitionTo(StatePtr nextState) {
        if (currentState) {
            currentState->onExit();
        }
        currentState = nextState;
        currentState->onEnter();
    }
};
}  // namespace os
/* --------- 语法糖 --------- */

namespace detail {
/**
 * @class TransitionBuilder
 * @brief 临时的构建器对象， state<S> + event<E> 会返回一个临时构建器
 */
class TransitionBuilder {
    os::State& state_;
    os::Event& event_;

public:
    TransitionBuilder(os::State& state, os::Event& event) : state_(state), event_(event) {}

    /**
     * @brief 为 state<S1> + event<E> >> state<S2> 重载
     */
    [[nodiscard]] StateCollector operator>>(os::State& targetState) const {
        // 对于直接转移，创建一个总是返回true的lambda作为条件的Branch
        state_.addLogic(event_, os::branch(Case{[]() noexcept { return true; }, targetState}));
        return StateCollector{{&state_, &targetState}};
    }

    /**
     * @brief 为 state<S> + event<E> >> BRANCH(...) 重载
     */
    [[nodiscard]] StateCollector operator>>(std::unique_ptr<Branch> branch) const {
        auto branch_states = branch->get_branch_state();
        std::vector<os::State*> new_states;
        new_states.reserve(branch_states.size() + 1);
        new_states.push_back(&state_);
        new_states.insert(new_states.end(), branch_states.begin(), branch_states.end());

        state_.addLogic(event_, std::move(branch));
        return StateCollector{std::move(new_states)};
    }
};

}  // namespace detail
/**
 * @brief state<S> + event<E>
 */
[[nodiscard]] inline detail::TransitionBuilder operator+(os::State& state, os::Event& event) {
    return detail::TransitionBuilder{state, event};
}

/**
 * @brief *state<S> : 仅声明一个状态的存在，确保它被绑定到状态机
 */
[[nodiscard]] inline detail::StateCollector operator*(os::State& state) { return detail::StateCollector{{&state}}; }

/**
 * @brief condition >> state: 创建一个带条件的Case
 */
template <typename Guard,
          // SFINAE 约束，避免与TransitionBuilder的成员operator>>冲突
          typename = std::enable_if_t<!std::is_same_v<std::decay_t<Guard>, detail::TransitionBuilder>>>
[[nodiscard]] auto operator>>(Guard&& guard, os::State& toState) {
    static_assert(std::is_invocable_r_v<bool, Guard>, "Guard must be a callable that returns bool.");
    return detail::Case{std::forward<Guard>(guard), toState};
}
#define CASE(condition) ([]() noexcept -> bool { return (condition); })
#define BRANCH os::branch
#define STATE(name) os::state<name>()
#define EVENT(name) os::event<name>()

#endif  // FSM_H
